import os
import numpy as np
import pandas as pd


class XlsxLogger:


    def __init__(self, xlogger_dir="", **kwargs):

        self.save_dir = None
        self.need_save = False
        if xlogger_dir == "":
            return

        self.save_dir = xlogger_dir
        os.makedirs(self.save_dir, exist_ok=True)

        self.table = None
        self.cnt = 0


    def record_scalar(self, content, **kwargs):
        
        if self.save_dir is None:
            return

        if self.table is None:
            self.table = pd.DataFrame()
            self.need_save = True

        for k in content:
            if type(content[k]) == torch.Tensor:
                content[k] = content[k].numpy().tolist()

        cur_tb = pd.DataFrame(content)
        self.table = self.table.append(cur_tb, ignore_index=True)


    def save(self):
        
        if not self.need_save:
            return

        self.table.to_excel(
            os.path.join(self.save_dir, f"xlsx_logger{self.cnt}.xlsx"), 
            sheet_name="Sheet1")
        self.cnt += 1
        self.table = None
        self.need_save = False


if __name__ == "__main__":
    import torch

    path = "E:/workspace/SOMatch/tmp/xlsx"
    xl = XlsxLogger(path)

    xl.save()

    des = {
        "max" : torch.randint(0, 100, (10,)), 
        "avg" : torch.randint(0, 100, (10,)), 
        "l2l" : torch.randint(0, 100, (10,)), 
    }

    xl.record_scalar(des)
    xl.save()
    xl.record_scalar(des)
    xl.record_scalar(des)
    xl.record_scalar(des)
    xl.save()


